import cv2
import numpy as np
import os
from os.path import exists
from imutils import paths
import pickle
from tqdm import tqdm
import logging

# 配置logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 定义函数get_size,获取文件大小,返回文件大小的兆字节的浮点数表示
def get_size(file):
    return os.path.getsize(file) / (1024 * 1024.0)

# 定义函数createXY,用于从图像创建特征(X)和标签(y)
def createXY(train_folder, dest_folder, method='flat', batch_size=64):
    x_file_path = os.path.join(dest_folder, "X.pkl")
    y_file_path = os.path.join(dest_folder, "y.pkl")

    if exists(x_file_path) and exists(y_file_path):
        logging.info("X和y已经存在，直接读取")
        with open(x_file_path, "rb") as f:
            X = pickle.load(f)
        with open(y_file_path, "rb") as f:
            y = pickle.load(f)
    else:
        logging.info("读取所有图像，生成X和y")
        image_paths = list(paths.list_images(train_folder))

        X = []
        y = []

        num_batches = len(image_paths) // batch_size + (1 if len(image_paths) % batch_size else 0)

        for idx in tqdm(range(num_batches), desc="读取图像"):
            batch_images = []
            batch_labels = []

            start = idx * batch_size
            end = min((idx + 1) * batch_size, len(image_paths))

            for i in range(start, end):
                image_path = image_paths[i]
                img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)  # 修复了这里的变量名错误
                img = cv2.resize(img, (32, 32))
                batch_images.append(img)
                label = image_path.split(os.path.sep)[-1].split('.')[0]
                label = 1 if label == 'dog' else 0
                batch_labels.append(label)

            batch_images = np.array(batch_images)
            batch_pixels = batch_images.reshape(-1, 32 * 32)  # 将图像数据转换为二维数组

            X.extend(batch_pixels)
            y.extend(batch_labels)

        X = np.array(X)
        y = np.array(y)

        with open(x_file_path, "wb") as f:
            pickle.dump(X, f)
        with open(y_file_path, "wb") as f:
            pickle.dump(y, f)

    return X, y

# 主函数
def main():
    train_folder = r"C:\Users\86187\Desktop\cat_dog_data\data\test"
    dest_folder = "."
    method = 'flat'
    batch_size = 64

    X, y = createXY(train_folder, dest_folder, method, batch_size)
    logging.info(f"X.shape: {X.shape}")
    logging.info(f"y.shape: {y.shape}")

if __name__ == '__main__':
    main()